import os
from copy import deepcopy

from process_v18 import messages_to_function_call
from util.query_openai_api import LLMClient

# from util.query_qwen_vllm import LLMClient

if __name__ == '__main__':
    import dotenv
    import json
    from concurrent.futures import ThreadPoolExecutor, as_completed

    dotenv.load_dotenv()
    llm = LLMClient(
        model=os.getenv("OPENAI_MODEL"),
        base_url=os.getenv("OPENAI_BASE_URL"),
        api_key=os.getenv("OPENAI_API_KEY"),

    )
    ds_llm = LLMClient(
        model=os.getenv("DS_MODEL"),
        api_key=os.getenv("DS_API_KEY"),
        base_url=os.getenv("DS_BASE_URL")
    )


    with open("res/半手工数据集.jsonl", "r", encoding="utf-8") as f:
        test_date = json.load(f)

    print("测试集载入完成: ", len(test_date))

    results_data = deepcopy(test_date)


    def process_data(i):
        print("正在生成第", i, "个数据")
        try:
            targets =messages_to_function_call(llm, test_date[i]["user_messages"], test_date[i]['apis'],
                                           llm_client_2=ds_llm)
        except Exception as e:
            print(f"处理第 {i} 个数据时发生错误: {e.with_traceback()}")
            targets = [[] for _ in range(len(test_date[i]["user_messages"]))]
        results_data[i]['targets']=targets
        print(targets)
        return targets

    with ThreadPoolExecutor(max_workers=50) as executor:
        future_to_index = {executor.submit(process_data, i): i for i in range(len(test_date))}
        for future in as_completed(future_to_index, timeout=600000):  # 设置超时为120秒
            index = future_to_index[future]
            try:
                result = future.result()
                # dataset.extend(result)
            except Exception as e:
                print(f"处理第 {index} 个数据时发生错误: {e}")

    print("完成: ")

    for i in range(len(results_data)):
        try:
            assert len(results_data[i]['targets']) == len(test_date[i]['user_messages'])
            assert results_data[i]['id'] == test_date[i]['id']
        except Exception as e:
            print(f"<UNK> {i}, {test_date[i]['id']} <UNK>",e)

    with open("result/test_4o.jsonl", "w", encoding="utf-8") as f:
        for l in results_data:
            # ，每个json对象一行
            s = json.dumps(l, ensure_ascii=False)
            f.writelines(s + "\n")
